﻿using System;
using System.Collections.Generic;
using System.Text;
using ASquared;
using ASquared.ModelStatistics;

namespace ASquared.ModelOptimization
{
    /// <summary>An interface defining a function with a single variable.</summary>
    public interface ISingleVariableFunction
    {
        /// <summary>Gets the derivative of the function with respect to its variable.</summary>
        /// <returns>Returns the derivative of the function.</returns>
        ISingleVariableFunction GetDeriv();
        /// <summary>Gets the function value evaluated at an input value.</summary>
        /// <param name="uVal">The value at which the function will be evaluated.</param>
        /// <returns>Returns the function value evaluated at uVal.</returns>
        Double GetValue(Double xVal);
        /// <summary>String representation of the function.</summary>
        /// <returns>Returns the string representation of the function.</returns>
        string ToString();
    }

    // Overall function 
    #region SingleVariableFunction

    public class SingleVariableFunction : ISingleVariableFunction
    {
        #region Instance variables
        
        public ISingleVariableFunction fxn;

        #endregion 

        #region Constructor

        public SingleVariableFunction()
        {
            this.fxn = new PowerFxn(1); 
        }
        public SingleVariableFunction(ISingleVariableFunction f)
        {
            this.fxn = f;
        }

        #endregion 

        #region Interface methods

        public virtual ISingleVariableFunction GetDeriv()
        {
            return this.fxn.GetDeriv();
        }
        public virtual Double GetValue(Double xVal)
        {
            return this.fxn.GetValue(xVal);
        }
        public override string ToString()
        {
            return this.fxn.ToString();
        }

        #endregion
        #region Generic methods

        public SingleVariableFunction Add(SingleVariableFunction f)
        {
            return this + f;
        }
        public SingleVariableFunction Subtract(SingleVariableFunction f)
        {
            return this - f;
        }
        public SingleVariableFunction Multiply(SingleVariableFunction f)
        {
            return this * f;
        }
        public SingleVariableFunction DivideBy(SingleVariableFunction f)
        {
            return this / f;
        }
        public SingleVariableFunction Pow(Double p)
        {
            return new SingleVariableFunction(new PowerFxn(p, this.fxn));
        }

        #endregion 
        #region Operators

        // Basic functions
        public static implicit operator SingleVariableFunction(AddFxn f)
        {
            return new SingleVariableFunction(f);
        }
        public static implicit operator SingleVariableFunction(MultiplyFxn f)
        {
            return new SingleVariableFunction(f);
        }
        public static implicit operator SingleVariableFunction(ConstantFxn f)
        {
            return new SingleVariableFunction(f);
        }
        public static implicit operator SingleVariableFunction(PowerFxn f)
        {
            return new SingleVariableFunction(f);
        }

        // Constant functions
        public static implicit operator SingleVariableFunction(float val)
        {
            return new ConstantFxn(val);
        }
        public static implicit operator SingleVariableFunction(Double val)
        {
            return new ConstantFxn(val);
        }
        public static implicit operator SingleVariableFunction(short val)
        {
            return new ConstantFxn(val);
        }
        public static implicit operator SingleVariableFunction(Int32 val)
        {
            return new ConstantFxn(val); 
        }
        public static implicit operator SingleVariableFunction(long val)
        {
            return new ConstantFxn(val);
        }

        // Advanced functions 
        public static implicit operator SingleVariableFunction(LogFxn f)
        {
            return new SingleVariableFunction(f); 
        }
        public static implicit operator SingleVariableFunction(PolynomialFxn f)
        {
            return new SingleVariableFunction(f); 
        }

        // String
        public static implicit operator string(SingleVariableFunction f)
        {
            if (f == null) return ""; 
            return f.ToString(); 
        }
        public static string operator +(string s, SingleVariableFunction f)
        {
            if (f == null) return s;
            return s + f.ToString(); 
        }
        public static string operator +(SingleVariableFunction f, string s)
        {
            if (f == null) return s;
            return f.ToString() + s; 
        }

        // Basic operators
        public static SingleVariableFunction operator +(SingleVariableFunction f, SingleVariableFunction g)
        {
            return new SingleVariableFunction(new AddFxn(f, g));
        }
        public static SingleVariableFunction operator -(SingleVariableFunction f, SingleVariableFunction g)
        {
            return new SingleVariableFunction(new AddFxn(f, new MultiplyFxn(-1, g)));
        }
        public static SingleVariableFunction operator *(SingleVariableFunction f, SingleVariableFunction g)
        {
            return new SingleVariableFunction(new MultiplyFxn(f, g));
        }
        public static SingleVariableFunction operator /(SingleVariableFunction f, SingleVariableFunction g)
        {
            return new SingleVariableFunction(new MultiplyFxn(f, new PowerFxn(-1, g)));
        }

        #endregion 
    }

    #endregion 

    // Base functions
    #region Add two functions

    /// <summary>Defines a function that is two other functions multiplied together f(x) + g(x).</summary>
    public class AddFxn : ISingleVariableFunction
    {
        public ISingleVariableFunction f;
        public ISingleVariableFunction g;

        // Properties
        public ConstantFxn ConstantPortion
        {
            get
            {
                if (this.f is ConstantFxn && this.g is ConstantFxn)
                    return ((ConstantFxn)this.f + (ConstantFxn)this.g);
                else if (this.f is ConstantFxn)
                    return (ConstantFxn)this.g;
                else if (this.g is ConstantFxn)
                    return (ConstantFxn)this.f;
                else
                    return null;
            }
        }
        public ISingleVariableFunction NonConstantPortion
        {
            get
            {
                ConstantFxn c = this.ConstantPortion;
                if (c == null)
                    return this;
                else if (c == this.f)
                    return this.g;
                else if (c == this.g)
                    return this.f;
                else
                    return null;
            }
        }

        // Constructor
        public AddFxn(ISingleVariableFunction f, ISingleVariableFunction g)
        {
            this.f = f;
            this.g = g;
        }
        public AddFxn(Double f, ISingleVariableFunction g)
            : this(new ConstantFxn(f), g)
        {
        }

        // Methods
        public ISingleVariableFunction GetDeriv()
        {
            if (this.f is ConstantFxn)
                return this.g.GetDeriv();
            else if (this.g is ConstantFxn)
                return this.f.GetDeriv();
            else
                return new AddFxn(this.f.GetDeriv(), this.g.GetDeriv());
        }
        public Double GetValue(Double xVal)
        {
            return this.f.GetValue(xVal) + this.g.GetValue(xVal);
        }
        public override string ToString()
        {
            ISingleVariableFunction n = this.NonConstantPortion; 
            ConstantFxn c = this.ConstantPortion;
            string fstr = this.f.ToString();
            string gstr = this.g.ToString();

            if (n == null)
                return c.ToString(); 
            if (c != null && c.val == 0)
                return this.NonConstantPortion.ToString();
            return fstr + " + " + gstr;
        }
    }

    #endregion
    #region Multiply two functions

    /// <summary>Defines a function that is two other functions multiplied together f(x) * g(x).</summary>
    public class MultiplyFxn : ISingleVariableFunction
    {
        public ISingleVariableFunction f;
        public ISingleVariableFunction g;

        // Constructor
        public MultiplyFxn(ISingleVariableFunction f, ISingleVariableFunction g)
        {
            this.f = f;
            this.g = g;
        }
        public MultiplyFxn(Double f, ISingleVariableFunction g)
            : this(new ConstantFxn(f), g)
        {
        }

        // Methods
        public ISingleVariableFunction GetDeriv()
        {
            ISingleVariableFunction fder = this.f.GetDeriv();
            ISingleVariableFunction gder = this.g.GetDeriv();
            ISingleVariableFunction prod1 = new MultiplyFxn(this.f, gder);
            ISingleVariableFunction prod2 = new MultiplyFxn(this.g, fder);
            if (this.f.GetType() == typeof(ConstantFxn))
                return prod1;
            else if (this.g.GetType() == typeof(ConstantFxn))
                return prod2;
            else
                return new AddFxn(prod1, prod2);
        }
        public Double GetValue(Double xVal)
        {
            return this.f.GetValue(xVal) * this.g.GetValue(xVal);
        }
        public override string ToString()
        {
            string fstr = this.f.ToString();
            string gstr = this.g.ToString();
            if (this.f is AddFxn)
                fstr = "(" + fstr + ")";
            if (this.g is AddFxn)
                gstr = "(" + gstr + ")";
            return fstr + "*" + gstr;
        }
    }

    #endregion
    #region ConstantFxn value

    /// <summary>Defines a function that is a constant = val.</summary>
    public class ConstantFxn : ISingleVariableFunction
    {
        public Double val;

        // Constructor 
        public ConstantFxn(Double value)
        {
            this.val = value;
        }
        public static implicit operator ConstantFxn(Double value)
        {
            return new ConstantFxn(value);
        }

        // Methods
        public ISingleVariableFunction GetDeriv()
        {
            return new ConstantFxn(0);
        }
        public Double GetValue(Double xVal)
        {
            return val;
        }
        public Double GetValue()
        {
            return val;
        }
        public override string ToString()
        {
            return val.ToString();
        }

        // Operators 
        public static ConstantFxn operator +(ConstantFxn a, ConstantFxn b)
        {
            return new ConstantFxn(a.val + b.val); 
        }
        public static ConstantFxn operator -(ConstantFxn a, ConstantFxn b)
        {
            return new ConstantFxn(a.val - b.val); 
        }
        public static ConstantFxn operator *(ConstantFxn a, ConstantFxn b)
        {
            return new ConstantFxn(a.val * b.val); 
        }
        public static ConstantFxn operator /(ConstantFxn a, ConstantFxn b)
        {
            return new ConstantFxn(a.val / b.val); 
        }
        public ConstantFxn Pow(Double pow)
        {
            return Math.Pow(this.val, pow); 
        }
    }

    #endregion
    #region Power function

    /// <summary>Defines a function with the form f(x) ^ p.</summary>
    public class PowerFxn : ISingleVariableFunction
    {
        // instance variables and properties 
        public Double p;
        public ISingleVariableFunction f;

        // Constructor 
        public PowerFxn(Double p)
            : this(p, null)
        {
        }
        public PowerFxn(Double p, ISingleVariableFunction f)
        {
            this.p = p;
            this.f = f;
        }

        // Methods
        public ISingleVariableFunction GetDeriv()
        {
            if (this.p == 0) return new ConstantFxn(0);
            if (this.p == 1) return new ConstantFxn(1);
            PowerFxn newderiv = new PowerFxn(this.p - 1, this.f);
            ConstantFxn a = new ConstantFxn(this.p);
            MultiplyFxn first = new MultiplyFxn(a, newderiv);
            if (this.f == null)
                return first;
            else
                return new MultiplyFxn(first, this.f.GetDeriv());
        }
        public Double GetValue(Double xVal)
        {
            if (this.f == null)
                return Math.Pow(xVal, this.p);
            else
                return Math.Pow(this.f.GetValue(xVal), this.p);
        }
        public override string ToString()
        {
            if (this.f == null)
            {
                if (this.p == 0)
                    return "1";
                else if (this.p == 1)
                    return "x";
                else
                    return "x^" + this.p.ToString();
            }
            else
            {
                if (p == 0)
                    return "1";
                else if (p == 1)
                    return this.f.ToString();
                else if (this.f is ConstantFxn) 
                    return ((ConstantFxn)this.f).Pow(this.p).ToString(); 
                else if (this.f is PowerFxn)
                {
                    PowerFxn pow = (PowerFxn)this.f;
                    pow = new PowerFxn(pow.p * this.p, pow.f); 
                    return pow.ToString();
                }
                else 
                    return "(" + this.f.ToString() + ")^" + this.p.ToString(); 
            }
        }
    }

    #endregion

    // Example functions
    #region Log function

    /// <summary>Defines a log function with logBase as its base: log_logBase(f(x)).</summary>
    public class LogFxn : ISingleVariableFunction
    {
        // instance variables and properties 
        public Double logBase;
        public ISingleVariableFunction f;

        // Constructor 
        public LogFxn()
            : this(Math.Exp(1), null)
        {
        }
        public LogFxn(Double logBase)
            : this(logBase, null)
        {
        }
        public LogFxn(ISingleVariableFunction f)
            : this(Math.Exp(1), f)
        {
        }
        public LogFxn(Double logBase, ISingleVariableFunction f)
        {
            this.logBase = logBase;
            this.f = f;
        }

        // Methods
        public ISingleVariableFunction GetDeriv()
        {
            PowerFxn newderiv = new PowerFxn(-1);
            ConstantFxn a = new ConstantFxn(1 / Math.Log(this.logBase));
            MultiplyFxn first = new MultiplyFxn(a, newderiv);
            if (this.f == null)
                return first;
            else
                return new MultiplyFxn(first, this.f.GetDeriv());
        }
        public Double GetValue(Double xVal)
        {
            if (this.f == null)
                return Math.Log(xVal, logBase);
            else
                return Math.Log(this.f.GetValue(xVal), logBase);
        }
        public override string ToString()
        {
            if (this.f == null)
            {
                if (logBase == Math.Exp(1))
                    return "ln(x)";
                else
                    return "log_" + logBase.ToString() + "(x)";
            }
            else
            {
                if (logBase == Math.Exp(1))
                    return "ln(" + this.f.ToString() + ")";
                else
                    return "log_" + logBase.ToString() + "(" + this.f.ToString() + ")";
            }
        }

    }

    #endregion
    #region PolynomialFxn function

    /// <summary>Defines a function with the form a0 + a1 * f1(x) + a2 * f2(x) ^ 2 + ... + aN * fN(x) ^ N.</summary>
    public class PolynomialFxn : ISingleVariableFunction
    {
        // instance variables and properties 
        private Double[] a;
        private ISingleVariableFunction[] f;
        private ISingleVariableFunction total;
        /// <summary>Gets the array of coefficients for this instance.</summary>
        public Double[] Coefficients
        {
            get { return this.a; }
        }
        /// <summary>Gets the array of functions for this instance.</summary>
        public ISingleVariableFunction[] Functions
        {
            get { return this.f; }
        }

        // Constructor 
        public PolynomialFxn(Double[] a)
            : this(a, null)
        {
        }
        public PolynomialFxn(Double[] a, ISingleVariableFunction[] f)
        {
            // error checks
            if (f != null && f.Length + 1 != a.Length && f.Length != 1) throw new Exception("The function ('f') array needs to have one less item than the coefficient ('a') array or just one item.");

            // coefficients
            this.a = a;

            // fill the functions array (with the same function if f.Length == 1) 
            if (f != null)
            {
                if (f.Length == 1)
                {
                    this.f = new ISingleVariableFunction[a.Length - 1];
                    for (Int32 i = 0; i + 1 < a.Length; i++)
                        this.f[i] = f[0];
                }
                else
                    this.f = f;
            }

            // create the polynomial
            this.total = this.CreatePolynomial();
        }

        // Methods
        private ISingleVariableFunction CreatePolynomial()
        {
            if (this.a.Length == 0) return new ConstantFxn(0);
            ConstantFxn a0 = this.a[0];
            if (this.a.Length == 1) return a0;
            ConstantFxn ai = 0.0;
            AddFxn tot = new AddFxn(a0, ai);
            for (Int32 i = 1; i < this.a.Length; i++)
            {
                ai = this.a[i];
                if (this.f == null)
                    tot = new AddFxn(tot, new MultiplyFxn(ai, new PowerFxn(i)));
                else
                    tot = new AddFxn(tot, new MultiplyFxn(ai, new PowerFxn(i, this.f[i - 1])));
            }
            return tot;
        }
        public ISingleVariableFunction GetDeriv()
        {
            return this.total.GetDeriv();
        }
        public Double GetValue(Double xVal)
        {
            return this.total.GetValue(xVal);
        }
        public override string ToString()
        {
            return this.total.ToString();
        }
        /// <summary>Gets the values from the polynomial with specified 'x' values.</summary>
        /// <param name="xVals">The 'x' values.</param>
        public Double[] GetValues(Double[] xVals)
        {
            Double[] polyVals = new Double[xVals.Length];
            for (Int32 i = 0; i < xVals.Length; i++)
                polyVals[i] = this.GetValue(xVals[i]);
            return polyVals; 
        }

        // Fitting polynomials 
        public static Int32 _maxPolyOrder = 10;
        public static PolynomialFxn Fit(Double[] xVals, Double[] yVals, Double targetR2)
        {
            ModelPerformance perf;
            return Fit(xVals, yVals, targetR2, out perf); 
        }
        public static PolynomialFxn Fit(Double[] xVals, Double[] yVals, Double targetR2, out ModelPerformance perf)
        {
            for (Int32 order = 0; order < _maxPolyOrder; order++)
            {
                PolynomialFxn poly = PolynomialFxn.Fit(xVals, yVals, order);
                perf = new ModelPerformance(yVals, poly.GetValues(xVals));
                if (perf.Rsquared() >= targetR2)
                    return poly;
            }
            throw new Exception("Could not match the target R² with any polynomial order up to " + _maxPolyOrder.ToString()); 
        }
        public static PolynomialFxn Fit(Double[] xVals, Double[] yVals, Int32 order)
        {
            if (xVals.Length != yVals.Length) throw new ArgumentException("xVals and yVals need to have the same number of elements in them."); 
            Matrix X = xVals; 
            Matrix y = yVals; 

            // Build columns (first column of 1's) 
            Matrix CurrCol = new Matrix(xVals.Length, 1.0);
            Matrix A = CurrCol; 
            for (Int32 j = 0; j < order; j++)
                A.AppendCols(CurrCol = CurrCol.ElementwiseMult(X)); 

            // Transpose 
            Matrix T = A.Transpose(); 

            // Least squares approx. fitting method
            Double[] a = (T * A).Inverse() * T * y;
            return new PolynomialFxn(a); 
        }
        public static string GoodnessOfFit(Double[] xVals, Double[] yVals, Double targetR2)
        {
            if (xVals.Length != yVals.Length) throw new Exception("The length in each of the input arrays must be the same.");
            PolynomialFxn p = PolynomialFxn.Fit(xVals, yVals, targetR2); 
            ModelPerformance perf = new ModelPerformance(yVals, p.GetValues(xVals));
            return "Goodness of fit for " + p.ToString() + ": \n" + perf.ToString();
        }
        public static string GoodnessOfFit(Double[] xVals, Double[] yVals, Int32 order)
        {
            if (xVals.Length != yVals.Length) throw new Exception("The length in each of the input arrays must be the same.");
            PolynomialFxn p = PolynomialFxn.Fit(xVals, yVals, order);
            ModelPerformance perf = new ModelPerformance(yVals, p.GetValues(xVals));
            return "Goodness of fit for " + p.ToString() + ": \n" + perf.ToString();
        }
        public static string GoodnessOfFit(Double[] xVals, Double[] yVals, PolynomialFxn p)
        {
            if (xVals.Length != yVals.Length) throw new Exception("The length in each of the input arrays must be the same."); 
            ModelPerformance perf = new ModelPerformance(yVals, p.GetValues(xVals));
            return "Goodness of fit for " + p.ToString() + ": \n" + perf.ToString();
        }
    }

    #endregion
}
